import os
from langchain import hub
# from langchain_deepseek import ChatDeepSeek
from langgraph.prebuilt import create_react_agent 
from langchain_community.tools.tavily_search import TavilySearchResults

# 设置TAVILY和DeepSeek API
os.environ["DEEPSEEK_API_KEY"] = ""
os.environ["TAVILY_API_KEY"] = ""

# 1. 从langchain的hub获取prompt的模板
prompt = hub.pull("zhang1career/react-agent-executor")
prompt.pretty_print()

# 2. 创建TavilySearchResults工具，设置最大结果数为1
tools = [TavilySearchResults(max_results = 1)]

# 3. 选择llm
# llm = ChatDeepSeek(
#     model="deepseek-chat",
#     api_key=os.getenv("DEEPSEEK_API_KEY"),
# )
from langchain_ollama import ChatOllama
llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)

# 4. 创建一个REACT代理执行器，使用指定的LLM和工具，并应用从Hub中获取的prompt 
agent_executor = create_react_agent(llm, tools, prompt=prompt)
 
import operator
from pydantic import BaseModel, Field
from typing import Annotated, List, Tuple, TypedDict,Literal
from langchain_core.prompts import ChatPromptTemplate
 
# 定义一个TypeDict类 PlanExecute, 用于存储输入、计划、过去的步骤和响应
class PlanExecute(TypedDict):
    input: str
    plan:List[str]
    past_steps:Annotated[List[Tuple], operator.add]
    response:str 

# 定义一个plan模型类， 用于描述未来要执行的计划
class Plan(BaseModel):
    """未来要执行的计划"""
    steps:List[str] = Field(
        description="需要执行的不同步骤，应该按顺序排列"
    )
 
# 创建一个计划生成的提示模板
planner_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """对于给定的目标，提出一个简单的逐步计划。这个计划应该包含独立的任务，如果正确执行将得出正确答案，不要添加任何多余的步骤，最后一步的结果应该是最终答案。确保每一步都有所有必要的信息 - 不要跳过步骤。"""
        ),
        ("placeholder","{messages}")
    ]
)

# | 管道操作符, LangChain的语法，表示将前一个组件的输出作为下一个组件的输入, 将提示模板连接到LLM
# 使用指定的提示词模板创建一个计划生成器，使用deepseek模型
# planner = planner_prompt|ChatDeepSeek(model="deepseek-chat",api_key=os.getenv("DEEPSEEK_API_KEY"),).with_structured_output(Plan)
planner = planner_prompt|llm.with_structured_output(Plan)

from typing import Union
 
# 定义一个响应模型类，用于描述用户的响应
class Response(BaseModel):
    """用户响应"""
    response: str
 
# 定义一个行为模型类，用于描述要执行的行为，该类继承自BaseModel
class Act(BaseModel):
    """要执行的行为"""
    action: Union[Response,Plan] = Field(
        description="要执行的行为， 如果要回应用户，使用Response，如果需要进一步使用工具获取答案，使用Plan。"
    )

# 重新规划步骤提示词 
replanner_prompt = ChatPromptTemplate.from_template(
    """对于给定的目标，提出一个简单的逐步计划。这个计划应该包含独立的任务，如果正确执行将得出正确答案，不要添加任何多余的步骤，最后一步的结果应该是最终答案。确保每一步都有所有必要的信息 - 
    你的目标是：
    {input}
    你的原计划是：
    {plan}
    你目前已完成的步骤是：
    {past_steps}
    相应的更新你的计划，如果不需要更多的步骤并且可以返回给用户，那么就这样响应。如果需要，填写计划。只添加仍然需要完成的步骤，不要返回已完成的步骤
    """
)   
 
# 使用指定的提示模板创建一个重新计划生成器，使用deepseek
# replanner = replanner_prompt|ChatDeepSeek(
#    model="deepseek-chat",
#     api_key=os.getenv("DEEPSEEK_API_KEY"),
# ).with_structured_output(Act)
replanner = replanner_prompt|llm.with_structured_output(Act)

import asyncio
# asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
 
from langgraph.graph import StateGraph, START

# 定义一个异步函数
async def main():
    # 定义一个异步函数，用于生成计划步骤
    async def plan_step(state: PlanExecute):
        plan = await planner.ainvoke({"messages": [("user", state["input"])]})
        return {"plan": plan.steps}
    
    # 定义一个异步函数执行步骤
    async def execute_step(state: PlanExecute):
        plan = state["plan"]
        plan_str = "\n".join(f"{i+1}.{step}" for i, step in enumerate(plan))
        task = plan[0]
        task_formatted = f"""对于一下计划:
         {plan_str}\n\n你的任务执行第{1}步，{task}.
           """
        agent_response = await agent_executor.ainvoke(
            {"messages":[("user", task_formatted)]}
        )
        return {"past_steps": state["past_steps"]+[(task, agent_response["messages"][-1].content)]}
    
 
    # 定义一个异步函数，用于重新计划步骤
    async def replan_step(state: PlanExecute):
        output = await replanner.ainvoke(state)
        if isinstance(output.action, Response):
            return {"response": output.action.response}
        else:
             return {"plan":output.action.steps}

    # 定义一个函数用于判断是否结束
    def should_end(state: PlanExecute)->Literal["agent","__end__"]:
        if("response" in state and state["response"]):
            return "__end__"
        else:
            return "agent"
 
    # 创建一个状态图，初始化PlanEexecute
    workflow = StateGraph(PlanExecute)
 
    # 添加计划节点
    workflow.add_node("planner",plan_step)
    # 添加执行步骤节点
    workflow.add_node("agent", execute_step)
    # 重新计划节点
    workflow.add_node("replan", replan_step)
 
    # 从开始到计划节点的边
    workflow.add_edge(START, "planner")
    # 设置从计划到代理节点的边
    workflow.add_edge("planner", "agent")
    # 设置从代理到重新计划节点的边
    workflow.add_edge("agent","replan")
 
    # 添加条件边，用于判断下一步操作
    workflow.add_conditional_edges(
        "replan",
        should_end,
    )
 
    app = workflow.compile()
    graph_png = app.get_graph().draw_mermaid_png()
    with open("agent_workflow.png", "wb") as f:
        f.write(graph_png)
    
    # 设置配置，递归限制为50次
    config = {"recursion_limit":50}
    inputs = {"input":"2025年WTT乒乓球男单冠军是是谁？请用中文回答"}
 
    async for event in app.astream(inputs, config=config):
        for k,v in event.items():
            if k!="__end__":
                print(v)
 
asyncio.run(main())
